# -*- coding: utf-8 -*-
"""
时区处理工具模块
目的：统一管理应用中的时间处理，确保时间数据的一致性
职责：提供时区转换、格式化等功能
"""
import pytz
from datetime import datetime
from typing import Optional, Union

# 定义北京时间时区
BEIJING_TZ = pytz.timezone('Asia/Shanghai')
UTC_TZ = pytz.UTC

def get_current_beijing_time() -> datetime:
    """
    获取当前北京时间
    目的：提供统一的北京时间获取方法

    Returns:
        datetime: 当前北京时间（带时区信息）
    """
    return datetime.now(BEIJING_TZ)

def convert_to_beijing_time(dt: Union[datetime, str]) -> datetime:
    """
    将时间转换为北京时间
    目的：确保所有时间数据都以北京时间存储和显示

    Args:
        dt: 需要转换的时间，可以是datetime对象或ISO格式字符串

    Returns:
        datetime: 转换后的北京时间
    """
    if isinstance(dt, str):
        dt = datetime.fromisoformat(dt.replace('Z', '+00:00'))

    if dt.tzinfo is None:
        # 如果时间对象没有时区信息，假设它是UTC时间
        dt = UTC_TZ.localize(dt)

    return dt.astimezone(BEIJING_TZ)

def format_beijing_time(dt: Optional[datetime], format_string: str = '%Y-%m-%d %H:%M:%S') -> Optional[str]:
    """
    格式化北京时间为字符串
    目的：提供统一的时间格式化方法

    Args:
        dt: 需要格式化的datetime对象
        format_string: 格式化字符串，默认格式为'YYYY-MM-DD HH:MM:SS'

    Returns:
        str: 格式化后的时间字符串，如果时间对象为None则返回None
    """
    if dt is None:
        return None

    # 确保时间有时区信息
    if dt.tzinfo is None:
        # 如果没有时区信息，假设它是北京时间
        dt = BEIJING_TZ.localize(dt)
    elif dt.tzinfo != BEIJING_TZ:
        # 如果不是北京时间，进行转换
        dt = dt.astimezone(BEIJING_TZ)

    return dt.strftime(format_string)

def get_beijing_today() -> datetime.date:
    """
    获取北京今天的日期
    目的：确保日期数据与北京时间一致

    Returns:
        date: 北京今天的日期
    """
    return get_current_beijing_time().date()

def ensure_beijing_time(dt: Optional[datetime]) -> Optional[datetime]:
    """
    确保datetime对象是北京时间
    目的：在存储到数据库前确保时间时区正确

    Args:
        dt: datetime对象，可能带或不带时区信息

    Returns:
        datetime: 带有时区信息的北京时间对象
    """
    if dt is None:
        return None

    if dt.tzinfo is None:
        # 无时区信息，假设为北京时间
        return BEIJING_TZ.localize(dt)
    else:
        # 有时区信息，转换为北京时间
        return dt.astimezone(BEIJING_TZ)

def utc_to_beijing(utc_dt: datetime) -> datetime:
    """
    将UTC时间转换为北京时间
    目的：处理从数据库读取的UTC时间

    Args:
        utc_dt: UTC时间

    Returns:
        datetime: 北京时间
    """
    if utc_dt.tzinfo is None:
        utc_dt = UTC_TZ.localize(utc_dt)
    return utc_dt.astimezone(BEIJING_TZ)

def beijing_to_utc(beijing_dt: datetime) -> datetime:
    """
    将北京时间转换为UTC时间
    目的：在存储到数据库前转换时间

    Args:
        beijing_dt: 北京时间

    Returns:
        datetime: UTC时间
    """
    if beijing_dt.tzinfo is None:
        beijing_dt = BEIJING_TZ.localize(beijing_dt)
    return beijing_dt.astimezone(UTC_TZ)